import os
import sys
import json
from datetime import datetime

from langchain_community.document_loaders import DirectoryLoader, TextLoader

from agents.architect import ArchitectAgent
from agents.programmer import ProgrammerAgent 
from agents.reviewer import ReviewerAgent
from agents.coordinator import CoordinatorAgent
from agents.optimizer import OptimizerAgent

from utils.embeddings import setup_rag,create_single_file_vector_store
from utils.callgraph import CallGraphBuilder
from config import PROJECT_PATH, TARGET_FILE, OUTPUT_PATH, EXCLUDED_DIRS, EXCLUDED_EXTENSIONS, MAX_ITERATIONS, OLLAMA_BASE_URL, LLM_MODEL
from utils.embeddings import get_file_vector
def main():
    """主程序入口点"""
    # 使用配置文件中的设置
    project_path = PROJECT_PATH
    target_file = TARGET_FILE
    output_path = OUTPUT_PATH
    max_iterations = MAX_ITERATIONS
    
    # 验证项目路径
    if not os.path.exists(project_path) or not os.path.isdir(project_path):
        print(f"错误: 项目路径 {project_path} 不存在或不是目录")
        return
    
    # 验证目标文件
    if target_file:
        # 如果目标文件是绝对路径，转换为相对于项目路径的路径
        if os.path.isabs(target_file):
            if target_file.startswith(project_path):
                # 将绝对路径转换为相对路径
                target_file = os.path.relpath(target_file, project_path)
            else:
                print(f"错误: 目标文件 {target_file} 不在项目路径 {project_path} 内")
                return
        
        full_target_path = os.path.join(project_path, target_file)
        if not os.path.exists(full_target_path) or not os.path.isfile(full_target_path):
            print(f"错误: 目标文件 {full_target_path} 不存在或不是文件")
            return
    
    # 确保输出目录存在
    os.makedirs(output_path, exist_ok=True)
    
    # 显示配置信息
    print("\n" + "="*50)
    print("📋 代码文档生成器配置")
    print("="*50)
    print(f"📁 项目路径: {project_path}")
    print(f"📄 目标文件: {target_file if target_file else '未指定 (分析整个项目)'}")
    print(f"📂 输出目录: {output_path}")
    print(f"🔄 最大迭代次数: {max_iterations}")
    print(f"🤖 Ollama服务: {OLLAMA_BASE_URL}")
    print(f"🧠 模型: {LLM_MODEL}")
    print("="*50 + "\n")
    
    # 无论是否有目标文件，都扫描整个项目构建调用关系树
    print("🔍 正在构建整个项目的调用关系树...")
    call_graph_builder = CallGraphBuilder(project_path)
    call_graph = call_graph_builder.build_full_graph()
    
    # 将调用关系保存为JSON
    project_name = os.path.basename(os.path.normpath(project_path))
    call_graph_file = os.path.join(output_path, f"{project_name}_call_graph.json")
    #call_graph_file = os.path.join(output_path, 'call_graph.json')
    with open(call_graph_file, 'w', encoding='utf-8') as f:
        json.dump(call_graph, f, indent=2)
    print(f"✅ 调用关系树已保存至: {call_graph_file}")
    
    # 为整个项目设置RAG系统，提供完整上下文
    print("📚 正在为整个项目建立RAG索引...")
    vector_store = setup_rag(project_path)
    print("✅ RAG索引构建完成")
    
    print("\n" + "="*50)
    print("🚀 初始化 Agents...")
    print("="*50)
    
    # 测试Ollama连接
    try:
        from langchain_community.llms import Ollama
        test_llm = Ollama(base_url=OLLAMA_BASE_URL, model=LLM_MODEL)
        test_response = test_llm("简短测试")
        print(f"✅ Ollama连接正常，使用模型: {LLM_MODEL}")
    except Exception as e:
        print(f"❌ Ollama连接失败: {str(e)}")
        print(f"请检查Ollama服务是否启动，以及模型 {LLM_MODEL} 是否已安装")
        return
    
    # 初始化各个Agent
    print("⏳ 初始化架构师 Agent...")
    architect_agent = ArchitectAgent(vector_store, project_path, call_graph)
    
    print("⏳ 初始化程序员 Agent...")
    # 在main.py中需要修改的部分（大约第98行附近）
    file_vector = None
    if target_file:
        full_target_path = os.path.join(project_path, target_file)
        if os.path.exists(full_target_path):
            file_vector = get_file_vector(full_target_path)
        else:
            print(f"⚠️ 警告: 目标文件 {full_target_path} 不存在，将继续但不使用文件向量")

    single_store = create_single_file_vector_store(full_target_path)
    programmer_agent = ProgrammerAgent(vector_store, project_path, call_graph, file_vector)
    
    print("⏳ 初始化评审员 Agent...")
    reviewer_agent = ReviewerAgent(single_store, full_target_path)

    print("⏳ 初始化优化员 Agent...")
    optimizer_agent = OptimizerAgent() 
    
    print("⏳ 初始化协调器 Agent...")
    coordinator_agent = CoordinatorAgent(architect_agent, programmer_agent, reviewer_agent, optimizer_agent)
    

    # 确定要分析的文件
    if target_file:
        # 单文件模式 - 只分析指定的文件，但利用整个项目的上下文
        file_paths = [target_file]
        print(f"🎯 单文件模式: 基于整个项目上下文，专注分析文件: {target_file}")
    else:
        # 多文件模式 - 扫描整个项目找出所有需要分析的文件
        file_paths = []
        for root, dirs, files in os.walk(project_path):
            # 过滤掉不需要的目录
            dirs[:] = [d for d in dirs if not any(excluded in d for excluded in EXCLUDED_DIRS)]
            
            for file in files:
                # 排除不需要分析的文件类型
                if not any(file.endswith(ext) for ext in EXCLUDED_EXTENSIONS):
                    rel_path = os.path.relpath(os.path.join(root, file), project_path)
                    file_paths.append(rel_path)
        
        print(f"📋 多文件模式: 将分析项目中的 {len(file_paths)} 个文件")
    
    if not file_paths:
        print("❌ 错误: 没有找到可分析的文件")
        return
    
    try:
        # 开始生成文档
        print("\n" + "="*50)
        print("🚀 开始生成文档...")
        print("="*50)
        result = coordinator_agent.generate_documentation(file_paths, max_iterations)
        
        # 为文件名添加前缀
        file_prefix = os.path.splitext(os.path.basename(target_file))[0] if target_file else "documentation"
        
        # 保存最终文档
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        doc_file = os.path.join(output_path, f'{file_prefix}_{timestamp}.md')
        with open(doc_file, 'w', encoding='utf-8') as f:
            f.write(result['documentation'])
        print(f"📄 文档已保存至: {doc_file}")
        
        # 保存详细步骤日志
        # log_file = os.path.join(output_path, f'{file_prefix}_steps_log_{timestamp}.json')
        # with open(log_file, 'w', encoding='utf-8') as f:
        #     json.dump(result['steps_log'], f, indent=2)
        # print(f"📊 步骤日志已保存至: {log_file}")
        
        print("\n" + "="*50)
        print("✨ 文档生成完成!")
        print("="*50)
        
    except KeyboardInterrupt:
        print("\n\n⚠️ 用户中断了操作")
        sys.exit(1)
    except Exception as e:
        print("\n" + "="*50)
        print(f"❌ 错误: 文档生成过程中发生异常:")
        print(f"{str(e)}")
        import traceback
        print(traceback.format_exc())
        print("="*50)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n⚠️ 用户中断了操作")
        sys.exit(1)
    except Exception as e:
        print(f"❌ 程序执行失败: {str(e)}")
        import traceback
        print(traceback.format_exc())
